import io
import json
from flask import Flask, request, Response
import argparse
import logging

logging.basicConfig(format='%(asctime)s.%(msecs)03d [%(levelname)s] %(message)s',
                    datefmt='## %Y-%m-%d %H:%M:%S')

logging.getLogger().setLevel(logging.INFO)
logger = logging.getLogger()

from server.server_model_hf import *

APP = Flask(__name__)

@APP.route('/infer', methods=['POST'])
def infer_image_func():
    ret = {
        'now_context': None,
        'full_context': None,
    }
    picture = request.files.get('picture')
    question = request.form.get('question')
    context = request.form.get('context', '')
    ocr_tokens = request.form.get('ocr_tokens', '')
    if(question is None or picture is None):
        return  "输入有误", 500
    kwarg = dict(
        do_sample = str2bool(request.form.get('sampling','true')),
        num_beams = int(request.form.get('num_beams','1')),
        top_p = float(request.form.get('top_p','0.9')),
        temperature = float(request.form.get('temperature','1')),
        num_return_sequences = int(request.form.get('num_captions','1')),
        repetition_penalty = float(request.form.get('repetition_penalty','1.5')),
        length_penalty = float(request.form.get('length_penalty','1')),
        max_length = int(request.form.get('max_length','500')),
        max_new_tokens = int(request.form.get('max_new_tokens','250')),
    )
    deterministic = str2bool(request.form.get('deterministic','false'))
    context = truncate_context(context, kwarg['max_length'], kwarg['max_new_tokens'])
    ocr_tokens = truncate_ocr(ocr_tokens, kwarg['max_length'], kwarg['max_new_tokens'])
    question = ocr_tokens + question
    kwarg.pop('max_length')

    try:
        picture_byte_stream = io.BytesIO(picture.read())
        Image.open(picture_byte_stream).convert('RGB')
    except:
        return "输入图片文件有误", 500
            
    generator = service_object.generate_stream(picture_byte_stream, question, context, deterministic, kwarg)

    def stream(generator):
        generated_text = ''
        for idx, text in enumerate(generator):
            # print(f'{text}: {text.encode()}')
            generated_text += text
            yield json.dumps({"index": idx,
                              "now_context": text},
                              ensure_ascii=False,
            ) + '\n'
        yield json.dumps({"index": idx,
                        "full_context": generated_text},
                        ensure_ascii=False,
            )
    
    return Response(stream(generator), mimetype="text/event-stream", status=200)


def str2bool(s):
    if(s == 'true' or s == 'True'):
        return True
    else:
        return False

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Training")

    parser.add_argument("--ckpt-path", required=False, 
                        default='/home/MLLM/250M_13B_GRD_half', help="模型文件地址")
    parser.add_argument("--is_lora", required=False, default=True, help="是否有Lora")
    args = parser.parse_args()

    service_object = Server_Model(args)

    APP.run(host="0.0.0.0", port=8080, threaded=True)